#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from typing import List, Dict, Any

class UnitFormatter:
    """单元格式化器，负责格式化输出格式"""
    
    def __init__(self):
        """初始化格式化器"""
        pass
    
    def format_unit_for_training(self, unit: Dict) -> Dict:
        """将unit格式化为训练格式"""
        return {
            'id': unit.get('unit_id', ''),
            'pr_number': unit.get('pr_number', ''),
            'file_path': unit.get('file_path', ''),
            'context_before': unit.get('context_before', ''),
            'context_after': unit.get('context_after', ''),
            'remove_content': unit.get('remove_content', ''),
            'add_content': unit.get('add_content', ''),
            'change_type': unit.get('change_type', ''),
            'hunk_header': unit.get('hunk_header', '')
        }
    
    def format_unit_for_analysis(self, unit: Dict) -> Dict:
        """将unit格式化为分析格式"""
        return {
            'unit_id': unit.get('unit_id', ''),
            'pr_number': unit.get('pr_number', ''),
            'file_path': unit.get('file_path', ''),
            'context_before': unit.get('context_before', ''),
            'context_after': unit.get('context_after', ''),
            'remove_content': unit.get('remove_content', ''),
            'add_content': unit.get('add_content', ''),
            'change_type': unit.get('change_type', ''),
            'hunk_header': unit.get('hunk_header', ''),
            'content_length': len(unit.get('add_content', '')) + len(unit.get('remove_content', ''))
        }
    
    def format_units_batch(self, units: List[Dict], format_type: str = 'training') -> List[Dict]:
        """批量格式化units"""
        if format_type == 'training':
            return [self.format_unit_for_training(unit) for unit in units]
        elif format_type == 'analysis':
            return [self.format_unit_for_analysis(unit) for unit in units]
        else:
            return units
    
    def validate_unit(self, unit: Dict) -> bool:
        """验证unit格式是否正确"""
        required_fields = ['pr_number', 'file_path', 'add_content', 'remove_content', 'change_type']
        
        for field in required_fields:
            if field not in unit:
                return False
        
        if not unit['pr_number'] or not unit['file_path']:
            return False
        
        return True
    
    def filter_valid_units(self, units: List[Dict]) -> List[Dict]:
        """过滤出有效的units"""
        return [unit for unit in units if self.validate_unit(unit)] 